package org.shiqianghh.model;

import java.nio.file.Paths;
import java.util.Map;

import org.shiqianghh.tanslate.DenoiserTranslator;
import org.shiqianghh.tanslate.SpeakerEncoderTranslator;
import org.shiqianghh.tanslate.TacotronSTFT;
import org.shiqianghh.tanslate.TacotronTranslator;
import org.shiqianghh.tanslate.WaveGlowTranslator;

import ai.djl.Device;
import ai.djl.inference.Predictor;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;

public class ModelUtils {
	
	public static NDArray get_mel(NDArray  wav) throws Exception{ 
		 
		Criteria<NDArray, NDArray> tacotronSTFT =
	             Criteria.builder()
	                     .setTypes(NDArray.class, NDArray.class)
	                     .optTranslator(new TacotronSTFT())
	                     .optEngine("PyTorch")
	                     .optDevice(Device.cpu())
	                     .optModelPath(Paths.get("build/voice/tacotronSTFT.pt"))
	                     .build();
		Predictor<NDArray, NDArray> stftPredictor = ModelZoo.loadModel(tacotronSTFT).newPredictor(); 
		wav = wav.expandDims(0);  
		NDArray melspec = stftPredictor.predict(wav);
        //melspec = self.stft.mel_spectrogram(audio_norm)
        //melspec = torch.squeeze(melspec, 0)
		melspec = melspec.squeeze(0);
        return melspec;
	}
	public static NDArray embed(NDArray frames_batch) throws Exception{
		Criteria<NDArray, NDArray> speakerEncoder =
	             Criteria.builder()
	                     .setTypes(NDArray.class, NDArray.class)
	                     .optTranslator(new SpeakerEncoderTranslator())
	                     .optEngine("PyTorch")
	                     .optDevice(Device.cpu())
	                     .optModelPath(Paths.get("build/voice/speakerEncoder.pt"))
	                     .build();
		
		Predictor<NDArray, NDArray> speakerEncoderPredictor = ModelZoo.loadModel(speakerEncoder).newPredictor(); 
		NDArray embed = speakerEncoderPredictor.predict(frames_batch);
       
       return embed;
	}
	public static Map<String,NDArray> inference_chu(NDList frames_batch) throws Exception{
		Criteria<NDList, Map> speakerEncoder =
	             Criteria.builder()
	                     .setTypes(NDList.class, Map.class)
	                     .optTranslator(new TacotronTranslator())
	                     .optEngine("PyTorch")
	                     .optDevice(Device.cpu())
	                     .optModelPath(Paths.get("build/voice/tacotron2.pt"))
	                     .build();
		
		Predictor<NDList, Map> speakerEncoderPredictor = ModelZoo.loadModel(speakerEncoder).newPredictor(); 
		Map<String,NDArray> embed = speakerEncoderPredictor.predict(frames_batch);
       //System.out.println(embed);
		return embed;
	}
	public static NDArray denoiser(NDArray wav,NDManager manager) throws Exception{
		Criteria<NDList, NDArray> denoiserEncoder =
	             Criteria.builder()
	                     .setTypes(NDList.class, NDArray.class)
	                     .optTranslator(new DenoiserTranslator())
	                     .optEngine("PyTorch")
	                     .optDevice(Device.cpu())
	                     .optModelPath(Paths.get("build/voice/denoiser.pt"))
	                     .build();
		
		Predictor<NDList, NDArray> denoiserPredictor = ModelZoo.loadModel(denoiserEncoder).newPredictor();
		NDArray denoiser_strength = manager.create(1.0f);
		NDList input = new NDList(); 
		
		/*NDList dim = new NDList(); 
		dim.add(wav);
		NDArray adddim = NDArrays.stack(dim);*/
		
		input.add(wav); 
		input.add(denoiser_strength); 
		NDArray embed = denoiserPredictor.predict(input);
		return embed;
	}
	
	public static NDArray generate_wave(NDArray mels_postnet,NDManager manager ) throws Exception{
		
		Criteria<NDList, NDArray> waveGlowEncoder =
	             Criteria.builder()
	                     .setTypes(NDList.class, NDArray.class)
	                     .optTranslator(new WaveGlowTranslator())
	                     .optEngine("PyTorch")
	                     .optDevice(Device.cpu())
	                     .optModelPath(Paths.get("build/voice/waveGlow.pt"))
	                     .build();
		
		Predictor<NDList, NDArray> waveGlowPredictor = ModelZoo.loadModel(waveGlowEncoder).newPredictor();
		NDArray sigma = manager.create(1.0);
		NDList input = new NDList(); 
		input.add(mels_postnet);
		input.add(sigma); 
		NDArray embed = waveGlowPredictor.predict(input);
      
		return embed; 
	}
	
}
